#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
from scapy.all import *
from typing import Dict, List, Tuple, Optional
import numpy as np
from datetime import datetime

class VPNDetector:
    """VPN协议检测器主类"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.detection_results = {}
        self.packet_buffer = []
        self.max_buffer_size = 1000
        
    def analyze_packet(self, packet: Packet) -> Dict[str, float]:
        """
        分析单个数据包，返回各VPN协议的可能性得分
        """
        results = {}
        
        # 1. OpenVPN检测
        results['openvpn'] = self._detect_openvpn(packet)
        
        # 2. IPSec/IKEv2检测
        results['ipsec'] = self._detect_ipsec(packet)
        
        # 3. WireGuard检测
        results['wireguard'] = self._detect_wireguard(packet)
        
        # 4. SSTP检测
        results['sstp'] = self._detect_sstp(packet)
        
        # 5. L2TP检测
        results['l2tp'] = self._detect_l2tp(packet)
        
        # 6. PPTP检测
        results['pptp'] = self._detect_pptp(packet)
        
        # 7. Cisco VPN检测
        results['cisco_vpn'] = self._detect_cisco_vpn(packet)
        
        # 8. AnyConnect检测
        results['anyconnect'] = self._detect_anyconnect(packet)
        
        # 9. SoftEther检测
        results['softether'] = self._detect_softether(packet)
        
        # 10. Shadowsocks检测
        results['shadowsocks'] = self._detect_shadowsocks(packet)
        
        # 11. V2Ray检测
        results['v2ray'] = self._detect_v2ray(packet)
        
        # 12. Trojan检测
        results['trojan'] = self._detect_trojan(packet)
        
        # 13. Tor检测
        results['tor'] = self._detect_tor(packet)
        
        # 14. Socks5检测
        results['socks5'] = self._detect_socks5(packet)
        
        # 15. SSH Tunnel检测
        results['ssh_tunnel'] = self._detect_ssh_tunnel(packet)
        
        # 16. FortiGate SSL VPN检测
        results['fortigate'] = self._detect_fortigate(packet)
        
        # 17. Palo Alto GlobalProtect检测
        results['palo_alto'] = self._detect_palo_alto(packet)
        
        # 18. Pulse Secure检测
        results['pulse_secure'] = self._detect_pulse_secure(packet)
        
        # 19. NordLynx检测
        results['nordlynx'] = self._detect_nordlynx(packet)
        
        # 20. Tailscale检测
        results['tailscale'] = self._detect_tailscale(packet)
        
        return results
    
    def _detect_openvpn(self, packet: Packet) -> float:
        """OpenVPN协议检测"""
        score = 0.0
        
        # 检查端口1194
        if packet.haslayer(TCP) or packet.haslayer(UDP):
            if packet[TCP].dport == 1194 or packet[TCP].sport == 1194:
                score += 0.3
                
        # 检查TLS握手特征
        if packet.haslayer(TLS):
            score += 0.4
            
        # 检查OpenVPN特定头部
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'OpenVPN' in raw_data:
                score += 0.3
                
        return min(score, 1.0)
    
    def _detect_ipsec(self, packet: Packet) -> float:
        """IPSec/IKEv2协议检测"""
        score = 0.0
        
        # 检查ISAKMP端口
        if packet.haslayer(UDP):
            if packet[UDP].dport == 500 or packet[UDP].sport == 500:
                score += 0.3
            if packet[UDP].dport == 4500 or packet[UDP].sport == 4500:
                score += 0.3
                
        # 检查ISAKMP头部
        if packet.haslayer(ISAKMP):
            score += 0.4
            
        return min(score, 1.0)
    
    def _detect_wireguard(self, packet: Packet) -> float:
        """WireGuard协议检测"""
        score = 0.0
        
        # 检查默认端口
        if packet.haslayer(UDP):
            if packet[UDP].dport == 51820 or packet[UDP].sport == 51820:
                score += 0.3
                
        # 检查固定148字节握手包
        if packet.haslayer(UDP) and len(packet[UDP].payload) == 148:
            score += 0.4
            
        return min(score, 1.0)
    
    def _detect_sstp(self, packet: Packet) -> float:
        """SSTP协议检测"""
        score = 0.0
        
        # 检查HTTPS特征
        if packet.haslayer(TCP) and packet[TCP].dport == 443:
            score += 0.3
            
        # 检查SSL/TLS特征
        if packet.haslayer(TLS):
            score += 0.4
            
        return min(score, 1.0)
    
    def _detect_l2tp(self, packet: Packet) -> float:
        """L2TP协议检测"""
        score = 0.0
        
        # 检查L2TP端口
        if packet.haslayer(UDP):
            if packet[UDP].dport == 1701 or packet[UDP].sport == 1701:
                score += 0.4
                
        # 检查IPSec封装
        if packet.haslayer(ESP):
            score += 0.3
            
        return min(score, 1.0)
    
    def _detect_pptp(self, packet: Packet) -> float:
        """PPTP协议检测"""
        score = 0.0
        
        # 检查PPTP端口
        if packet.haslayer(TCP):
            if packet[TCP].dport == 1723 or packet[TCP].sport == 1723:
                score += 0.4
                
        # 检查GRE封装
        if packet.haslayer(GRE):
            score += 0.3
            
        return min(score, 1.0)
    
    def _detect_cisco_vpn(self, packet: Packet) -> float:
        """Cisco VPN协议检测"""
        score = 0.0
        
        # 检查DTLS特征
        if packet.haslayer(UDP) and packet[UDP].dport == 443:
            score += 0.3
            
        # 检查Cisco特定标识
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'Cisco' in raw_data:
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_anyconnect(self, packet: Packet) -> float:
        """AnyConnect协议检测"""
        score = 0.0
        
        # 检查CSTP协议头
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'CSTP' in raw_data:
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_softether(self, packet: Packet) -> float:
        """SoftEther协议检测"""
        score = 0.0
        
        # 检查HTTPS伪装
        if packet.haslayer(TCP) and packet[TCP].dport == 443:
            score += 0.3
            
        # 检查特定心跳包间隔
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'SoftEther' in raw_data:
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_shadowsocks(self, packet: Packet) -> float:
        """Shadowsocks协议检测"""
        score = 0.0
        
        # 检查固定头结构
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if len(raw_data) > 0 and raw_data[0] in [0x01, 0x03, 0x04]:
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_v2ray(self, packet: Packet) -> float:
        """V2Ray协议检测"""
        score = 0.0
        
        # 检查VMess协议ID
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if len(raw_data) > 16 and raw_data[0:16].hex().startswith('00000000'):
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_trojan(self, packet: Packet) -> float:
        """Trojan协议检测"""
        score = 0.0
        
        # 检查TLS特征
        if packet.haslayer(TLS):
            score += 0.3
            
        # 检查特洛伊协议头
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'TROJAN' in raw_data:
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_tor(self, packet: Packet) -> float:
        """Tor协议检测"""
        score = 0.0
        
        # 检查Cell包结构
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if len(raw_data) == 514:  # 标准Tor cell大小
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_socks5(self, packet: Packet) -> float:
        """Socks5协议检测"""
        score = 0.0
        
        # 检查默认端口
        if packet.haslayer(TCP):
            if packet[TCP].dport == 1080 or packet[TCP].sport == 1080:
                score += 0.3
                
        # 检查认证协商模式
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if raw_data[0] == 0x05:  # SOCKS5版本
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_ssh_tunnel(self, packet: Packet) -> float:
        """SSH Tunnel协议检测"""
        score = 0.0
        
        # 检查SSH版本协商
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'SSH-' in raw_data:
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_fortigate(self, packet: Packet) -> float:
        """FortiGate SSL VPN协议检测"""
        score = 0.0
        
        # 检查特定Cookie值
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'FGT' in raw_data:
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_palo_alto(self, packet: Packet) -> float:
        """Palo Alto GlobalProtect协议检测"""
        score = 0.0
        
        # 检查GP头标识
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'GlobalProtect' in raw_data:
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_pulse_secure(self, packet: Packet) -> float:
        """Pulse Secure协议检测"""
        score = 0.0
        
        # 检查DSID Cookie特征
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'DSID' in raw_data:
                score += 0.4
                
        return min(score, 1.0)
    
    def _detect_nordlynx(self, packet: Packet) -> float:
        """NordLynx协议检测"""
        score = 0.0
        
        # 检查WireGuard改进特征
        if packet.haslayer(UDP):
            if packet[UDP].dport == 51820 or packet[UDP].sport == 51820:
                score += 0.3
                
        return min(score, 1.0)
    
    def _detect_tailscale(self, packet: Packet) -> float:
        """Tailscale协议检测"""
        score = 0.0
        
        # 检查DERP协议特征
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            if b'DERP' in raw_data:
                score += 0.4
                
        return min(score, 1.0)
    
    def process_pcap(self, pcap_file: str) -> Dict[str, List[Dict]]:
        """
        处理pcap文件并返回检测结果
        """
        results = []
        try:
            packets = rdpcap(pcap_file)
            for packet in packets:
                detection_result = self.analyze_packet(packet)
                results.append({
                    'timestamp': datetime.now().isoformat(),
                    'protocol_scores': detection_result
                })
        except Exception as e:
            self.logger.error(f"Error processing pcap file: {str(e)}")
            
        return results 